import csv
import logging
import os
from typing import Optional

from langchain_community.document_loaders import PyPDFLoader
from pathlib import Path
from .generateGraphUtils import *
import sys

sys.path.append("..")

from openai import OpenAI


class GenerateGraph:
    @classmethod
    def generateGraph(cls, file_path) -> list[Optional[None]]:
        logging.basicConfig(level=logging.INFO)
        try:
            api_key = get_api_key()
            client = OpenAI(
                api_key=api_key,
                base_url="https://ark.cn-beijing.volces.com/api/v3",
            )
            logging.info("client is successfully initialized")

            pdf_file_path = Path(file_path)
            # 使用 PyPDFLoader 加载 PDF 文件
            loader = PyPDFLoader(pdf_file_path)
            documents = loader.load()
            # 调用documents2Dataframe函数
            df = documents2Dataframe(documents)
            # 调用df2Graph函数
            concept_list = df2Graph(df, client, "ep-20250227161813-t6vpw")
            # 调用graph2Df函数
            graph_df = graph2Df(concept_list)
            generate_graph(concept_list)
            return concept_list
        except Exception as e:
            logging.error("client failed")
            print(e)
            return None

    @classmethod
    def save(cls,triples_list:list[Optional[None]],graph_directory:str,file_name:str) -> str:
        if not os.path.exists(graph_directory):
            os.makedirs(graph_directory)
        graph_file_path = os.path.join(graph_directory,file_name)

        with open(graph_file_path, mode='w', newline='', encoding='utf-8') as file:
            writer = csv.DictWriter(file, fieldnames=['chunk_id', 'node_1', 'node_2', 'edge'])
            writer.writeheader()  # 写入表头
            writer.writerows(triples_list)  # 写入数据

        print(f"数据已保存到 {graph_file_path}")
        return graph_file_path

